"""
Interface definitions for continual learning methods.

Continual learning methods are algorithms that manage model training
across multiple tasks to prevent catastrophic forgetting.
"""

from abc import abstractmethod
from typing import Any, Dict, List, Optional, Union

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from .core.component import ComponentInterface

"""
Protocol definitions for task adaptation in continual learning.
"""

from typing import Protocol, runtime_checkable


@runtime_checkable
class TaskAdaptProtocol(Protocol):
    """
    Protocol for components that can adapt to different tasks.

    This protocol defines the interface that components must implement
    to support task switching in continual learning scenarios.
    Components implementing this protocol can be automatically discovered
    and coordinated by the task boundary callback system.
    """

    def set_task(self, task_id: int) -> None:
        """
        Configure the component for a specific task.

        This method is called when switching between tasks to ensure
        the component is properly configured for the target task.

        Args:
            task_id: Identifier for the task to configure

        Raises:
            ValueError: If task_id is invalid or not supported
        """
        ...


class MethodInterface(ComponentInterface):
    """
    Interface for continual learning methods.

    A method represents a continual learning algorithm that manages model
    training across multiple tasks to prevent catastrophic forgetting.
    Examples include EWC, replay-based methods, and parameter isolation approaches.
    """

    @abstractmethod
    def on_task_start(
        self, task_id: int, task_dataloader: Optional[DataLoader] = None
    ) -> None:
        """
        Called at the beginning of a task.

        Use this method to prepare the method for a new task,
        e.g., by initializing task-specific parameters or buffers.

        Args:
            task_id: Task identifier
            task_dataloader: Optional dataloader with task data
        """
        pass

    @abstractmethod
    def on_task_end(
        self, task_id: int, task_dataloader: Optional[DataLoader] = None
    ) -> None:
        """
        Called at the end of a task.

        Use this method to perform operations after a task has been learned,
        e.g., updating importance weights or storing exemplars.

        Args:
            task_id: Task identifier
            task_dataloader: Optional dataloader with task data
        """
        pass

    @abstractmethod
    def compute_loss(
        self,
        model: nn.Module,
        criterion: nn.Module,
        output: Any,
        targets: torch.Tensor,
        task_id: int,
        **kwargs,
    ) -> Dict[str, torch.Tensor]:
        """
        Compute the loss for the given outputs and targets.

        This method should be implemented by all methods to define
        how the loss is computed for continual learning.

        Args:
            model: The model being trained
            criterion: The base criterion (loss function)
            output: Model output (logits or dict with logits)
            targets: Target labels
            task_id: Task identifier
            **kwargs: Additional arguments

        Returns:
            Dictionary with loss values (total loss under key 'loss')
        """
        pass

    def get_parameters_to_optimize(
        self, model: nn.Module
    ) -> Union[List[nn.Parameter], List[Dict]]:
        """
        Get parameters that should be optimized for the current task.

        This method can be overridden to provide different parameters
        to the optimizer, e.g., for parameter-efficient fine-tuning.

        Args:
            model: The model being trained

        Returns:
            List of parameters or list of parameter groups
        """
        return list(model.parameters())

    def state_dict(self) -> Dict[str, Any]:
        """
        Get state dictionary for serialization.

        Returns:
            Dictionary containing method state
        """
        return {}

    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
        """
        Load state dictionary for deserialization.

        Args:
            state_dict: Dictionary containing method state
        """
        pass
